import pandas as pd

# 创建示例数据
data_user = pd.DataFrame({
    'user_id': [1, 1, 3, 2, 2, 1],
    'date': pd.to_datetime(['2022-01-01', '2022-01-02', '2022-01-03', '2022-01-01', '2022-01-02', '2022-01-03']),
    'behavior_type': ['4', '4', '4', '4', '4', '4']
})

# 定义计算时间差的函数
def calculate_time_diff(group):
    group['time_diff'] = group['date'].diff()
    #group['time_diff'].iloc[0] = pd.NaT
    return group

# 分组计算时间差，同时保持分组特性
data_day_buy = data_user[data_user.behavior_type == '4'].groupby('user_id').apply(calculate_time_diff)

print(data_day_buy[['user_id', 'date', 'time_diff']])
